"""ChromaDB服务包装器 - 提供统一的ChromaDB操作接口"""

import logging
from typing import Dict, List, Any, Optional
from langchain_community.vectorstores import Chroma

from app.core.vector_store import get_vector_store

logger = logging.getLogger(__name__)


class ChromaService:
    """ChromaDB服务类 - 包装现有的vector_store功能"""
    
    def __init__(self):
        self.vector_store = None
        self._initialize_vector_store()
    
    def _initialize_vector_store(self):
        """初始化向量存储"""
        try:
            self.vector_store = get_vector_store()
            logger.info("✅ [ChromaService] ChromaDB初始化成功")
        except Exception as e:
            logger.error(f"❌ [ChromaService] ChromaDB初始化失败: {e}")
            self.vector_store = None
    
    def _ensure_vector_store(self):
        """确保向量存储可用"""
        if self.vector_store is None:
            self._initialize_vector_store()
        return self.vector_store is not None
    
    async def search_knowledge(self, query: str, n_results: int = 5, 
                             filters: Optional[Dict] = None) -> List[Dict[str, Any]]:
        """搜索知识库"""
        if not self._ensure_vector_store():
            logger.error("❌ [ChromaService] 向量存储不可用")
            return []
        
        try:
            # 使用similarity_search方法
            results = self.vector_store.similarity_search(
                query=query,
                k=n_results
            )
            
            # 转换为标准格式
            formatted_results = []
            for i, doc in enumerate(results):
                formatted_results.append({
                    "id": f"knowledge_{i}",
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "score": getattr(doc, 'score', 0.0)
                })
            
            logger.info(f"🔍 [ChromaService] 知识搜索完成，返回 {len(formatted_results)} 条结果")
            return formatted_results
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 知识搜索失败: {e}")
            return []
    
    async def search_conversations(self, user_id: str, query: str = "", 
                                 n_results: int = 10) -> List[Dict[str, Any]]:
        """搜索用户对话"""
        if not self._ensure_vector_store():
            logger.error("❌ [ChromaService] 向量存储不可用")
            return []
        
        try:
            # 构建搜索查询
            search_query = query if query else f"user:{user_id}"
            
            # 使用similarity_search_with_score方法
            results = self.vector_store.similarity_search_with_score(
                query=search_query,
                k=n_results
            )
            
            # 转换为标准格式
            formatted_results = []
            for i, (doc, score) in enumerate(results):
                # 过滤用户相关的对话
                if doc.metadata.get('user_id') == user_id or not user_id:
                    formatted_results.append({
                        "id": f"conversation_{i}",
                        "content": doc.page_content,
                        "metadata": doc.metadata,
                        "score": score
                    })
            
            logger.info(f"🔍 [ChromaService] 对话搜索完成，返回 {len(formatted_results)} 条结果")
            return formatted_results
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 对话搜索失败: {e}")
            return []
    
    async def similarity_search(self, query_text: str, collection_name: str = None,
                              n_results: int = 5, filters: Optional[Dict] = None) -> List[Dict[str, Any]]:
        """相似性搜索"""
        if not self._ensure_vector_store():
            logger.error("❌ [ChromaService] 向量存储不可用")
            return []
        
        try:
            # 使用similarity_search_with_score方法
            results = self.vector_store.similarity_search_with_score(
                query=query_text,
                k=n_results
            )
            
            # 转换为标准格式
            formatted_results = []
            for i, (doc, score) in enumerate(results):
                # 应用过滤器（如果提供）
                if filters:
                    match = True
                    for key, value in filters.items():
                        if doc.metadata.get(key) != value:
                            match = False
                            break
                    if not match:
                        continue
                
                formatted_results.append({
                    "id": f"similarity_{i}",
                    "content": doc.page_content,
                    "metadata": doc.metadata,
                    "score": score
                })
            
            logger.info(f"🔍 [ChromaService] 相似性搜索完成，返回 {len(formatted_results)} 条结果")
            return formatted_results
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 相似性搜索失败: {e}")
            return []
    
    def add_documents(self, documents: List[str], metadatas: List[Dict] = None, 
                     ids: List[str] = None) -> bool:
        """添加文档到向量存储"""
        if not self._ensure_vector_store():
            logger.error("❌ [ChromaService] 向量存储不可用")
            return False
        
        try:
            self.vector_store.add_texts(
                texts=documents,
                metadatas=metadatas,
                ids=ids
            )
            logger.info(f"✅ [ChromaService] 成功添加 {len(documents)} 个文档")
            return True
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 添加文档失败: {e}")
            return False
    
    def delete_documents(self, ids: List[str]) -> bool:
        """删除文档"""
        if not self._ensure_vector_store():
            logger.error("❌ [ChromaService] 向量存储不可用")
            return False
        
        try:
            self.vector_store.delete(ids=ids)
            logger.info(f"✅ [ChromaService] 成功删除 {len(ids)} 个文档")
            return True
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 删除文档失败: {e}")
            return False
    
    def get_collection_info(self) -> Dict[str, Any]:
        """获取集合信息"""
        if not self._ensure_vector_store():
            return {"status": "unavailable"}
        
        try:
            # 获取集合统计信息
            collection = self.vector_store._collection
            count = collection.count()
            
            return {
                "status": "available",
                "document_count": count,
                "collection_name": getattr(collection, 'name', 'unknown')
            }
            
        except Exception as e:
            logger.error(f"❌ [ChromaService] 获取集合信息失败: {e}")
            return {"status": "error", "error": str(e)}
